#!/usr/bin/env python3
"""
Direct SQL benchmark using kwbase sql command (subprocess)
This represents the baseline performance for direct SQL writes without SDK overhead.
"""

import os
import sys
import time
import subprocess
import threading
import random
from datetime import datetime, timedelta

success_count = 0
error_count = 0
lock = threading.Lock()

def exec_sql(sql):
    """Execute SQL using kwbase sql command"""
    kwbase_path = os.getenv('KWBASE_PATH', '/home/go/src/gitee.com/kwbasedb/install/bin')
    host = os.getenv('KWDB_HOST', '127.0.0.1')
    port = os.getenv('KWDB_SQL_PORT', '26257')
    db = os.getenv('KWDB_DB', 'kwdb')

    cmd = [
        'bash', '-c',
        f'cd {kwbase_path} && export LD_LIBRARY_PATH=../lib && '
        f'./kwbase sql --insecure --host={host}:{port} --database={db} -e "{sql}" 2>&1'
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode == 0:
            return True, ""
        else:
            return False, result.stdout + result.stderr
    except Exception as e:
        return False, str(e)

def worker_thread(thread_id, total_per_thread, batch_size):
    """Worker thread for concurrent writes"""
    global success_count, error_count

    # Monotonic timestamp generation (same as C++ benchmark)
    base_timestamp = datetime(2024, 9, 30) + timedelta(seconds=thread_id * 10000)
    point_counter = 0

    written = 0
    while written < total_per_thread:
        batch_count = min(batch_size, total_per_thread - written)

        # Build batch INSERT
        values = []
        for i in range(batch_count):
            ts = base_timestamp + timedelta(microseconds=point_counter)
            host = f"bench-host-{thread_id % 4}"
            region = f"bench-region-{thread_id % 2}"
            usage = random.uniform(0.0, 100.0)
            temperature = random.uniform(0.0, 100.0)

            ts_str = ts.strftime('%Y-%m-%d %H:%M:%S.%f')
            values.append(f"('{ts_str}', '{host}', '{region}', {usage:.2f}, {temperature:.2f})")
            point_counter += 1

        sql = f"INSERT INTO sdk_bench (ts, host, region, usage, temperature) VALUES {','.join(values)};"

        success, err = exec_sql(sql)

        with lock:
            if success:
                success_count += batch_count
            else:
                error_count += batch_count
                print(f"[thread-{thread_id}] Error: {err[:100]}", file=sys.stderr)

        written += batch_count

def main():
    if len(sys.argv) < 4:
        print("Usage: python3 sql_benchmark.py <total_points> <batch_size> <concurrency>")
        print("Example: python3 sql_benchmark.py 100000 200 8")
        print()
        print("Environment variables:")
        print("  KWDB_HOST      - database host (default: 127.0.0.1)")
        print("  KWDB_SQL_PORT  - database port (default: 26257)")
        print("  KWDB_DB        - database name (default: kwdb)")
        print("  KWBASE_PATH    - path to kwbase binary (default: /home/go/src/gitee.com/kwbasedb/install/bin)")
        sys.exit(1)

    total_points = int(sys.argv[1])
    batch_size = int(sys.argv[2])
    concurrency = int(sys.argv[3])

    if total_points <= 0 or batch_size <= 0 or concurrency <= 0:
        print("Error: All parameters must be positive integers", file=sys.stderr)
        sys.exit(1)

    points_per_thread = total_points // concurrency
    actual_total = points_per_thread * concurrency

    print("=== Direct SQL Benchmark ===")
    print(f"Total points:  {actual_total}")
    print(f"Batch size:    {batch_size}")
    print(f"Concurrency:   {concurrency}")
    print(f"Per thread:    {points_per_thread}")
    print(f"KWBASE_PATH:   {os.getenv('KWBASE_PATH', '/home/go/src/gitee.com/kwbasedb/install/bin')}")
    print()

    start_time = time.time()

    threads = []
    for i in range(concurrency):
        t = threading.Thread(target=worker_thread, args=(i, points_per_thread, batch_size))
        t.start()
        threads.append(t)

    for t in threads:
        t.join()

    end_time = time.time()
    duration_sec = end_time - start_time

    qps = success_count / duration_sec if duration_sec > 0 else 0
    avg_latency = (duration_sec * 1000 * concurrency) / (success_count / batch_size) if success_count > 0 else 0

    print("=== Results ===")
    print(f"Duration:      {duration_sec:.2f} seconds")
    print(f"Success:       {success_count} points")
    print(f"Errors:        {error_count} points")
    print(f"QPS:           {qps:.2f} points/sec")
    print(f"Avg latency:   {avg_latency:.2f} ms/batch")
    print()

    if error_count > 0:
        print(f"⚠ Warning: {error_count} points failed to write")
        sys.exit(1)

    print("✓ Benchmark completed successfully")

if __name__ == '__main__':
    main()
